import scipy
from scipy.stats import linregress
import numpy
from numpy import asarray
from numpy import linspace
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot as PLT
from CalTorque import *
import os
import sys
from datetime import datetime
from scipy.interpolate import UnivariateSpline

# NO point in making this a class

class Velocity: # No point in this begin a class
    """ Calculate velocity of source before S5, and overall velocity up to S5. """
    def __init__(self, pos, time, run, rpm):
        self.pos = pos
        self.time = time
        self.run = run
        self.rpm = rpm
        self.instVel = []
        self.InstVelocity()
        self.velocity = self.IntegratedVelocity()
        self.velocityAtS5 = self.VelocityBeforeS5()
        
    def InstVelocity(self):
        """ Calculates velocity by multiplying the angular velocity data by the radius
            of the sprocket wheel """
        r = 2 * 2.54 # radius of sprocket wheel in cm TODO check this value
        for w in self.rpm:
            w = 2 * w * scipy.pi / 60.
            self.instVel.append(w * 2 * 2.54)
    
    def IntegratedVelocity(self):
        """ Fits a spline to position versus time data with scipy's UnivariateSpline
            Takes the derivative of spline to get integrated velocity of run """
        timeChop = []
        posChop = []
        velChop = []
        
        for i, p in enumerate(self.pos):
            ## TODO more robust way of doing this
            if p == self.pos[0] or p == self.pos[-1] or p >= 385: continue
            posChop.append(p)
            timeChop.append(self.time[i])
            velChop.append(self.instVel[i])
        #print timeChop
        #print posChop
        
        spl = UnivariateSpline(timeChop, posChop) # s=2 works
        ys = spl(timeChop)
        
        #print ys
        #print spl(timeChop)
        
        d = spl.derivative()
        yprime = d(timeChop)
        
        ### TODO fix, and chop time appropriately
        fig = PLT.figure(figsize = (15, 8), dpi = 150)
        box = [0.14, 0.14, 0.76, 0.76]
        ax = fig.add_axes(box)
        ax2 = ax.twinx()
        ax.set_ylabel('Position (cm)', color = 'blue')
        ax.set_xlabel('Time (s)')
        ax2.set_ylabel('Velocity (cm/s)', color = 'green')
        PLT.title('Integrated Velocity Test')
        ax.plot(self.time, self.pos, color = 'black', label = "Position versus time data")
        ax.plot(timeChop, posChop, '-', color = 'black', label = "Chopped Position vs Time data")
        ax.plot(timeChop, ys, '-', color = 'red', label = "Spline")
        ax2.plot(timeChop, velChop, color = 'green', label = "Angular velocity * radius of sprocket")
        ax2.plot(timeChop, yprime, '-', color = 'blue', label = "Derivative of Spline - Velocity")
        h, l = ax.get_legend_handles_labels()
        q, t = ax2.get_legend_handles_labels()
        ax.set_ylim([-25,400])
        ax.legend(h, l, 'upper left')
        ax2.legend(q, t, 'lower right')
        fig.savefig('vel/IntegratedVelocityTest2' + self.run + '.png')
        PLT.close()
        
        return yprime

    def AverageVelocity(self):
        """ Determines the average velocity of deployment"""
        velocity = self.IntegratedVelocity()
        #print "-> 1", velocity
        counter = 0
        velocitySum = 0
        for v in self.velocity:
            counter += 1
            velocitySum += v
        #print velocitySum / counter
        return velocitySum / counter

    def MaximumVelocity(self):
        """ Determines the maximum velocity"""
        for i, v in enumerate(self.velocity):
            if i == 0:
                maxVelocity = v
            if v > maxVelocity:
                maxVelocity = v
        return maxVelocity

    def MedianVelocity(self):
        """ Determines the median velocity"""
        velocity = self.velocity
        velocity.sort()
        i = len(velocity)/2.
        if len(velocity) % 2 == 0:
            return (velocity[i] + velocity[i+1]) /2.
        else:
            return velocity[i]
    
    def Velocity2(self):
        """ Accessor for integrated velocity """
        return self.velocity
    
    def VelocityAtS5(self):
        """ Accessor for velocity at S5 """
        return self.velocityAtS5

#x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-- Testing --x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x-x#
def main(argv = sys.argv):
    
    database = OpenCalTorqueDB("database.pkl")
    
    for thing in database:
        if thing.Direction() != 'in': continue
        if thing.Station() != 'S5': continue
        try:
            print thing.Filename()
            vel = Velocity(thing.NewPosition(), thing.Time(), "".join(str(thing.Datetime()).split(' ')[0].split('-')), thing.rpm)
            vel.Velocity2()
            #print vel.AverageVelocity()
            #print vel.MaximumVelocity()
            #print vel.MedianVelocity()
        except:
            print "File could not be processed: ", thing.Filename()
    
    

if __name__ == "__main__":
   main(sys.argv)